import pykd

breaklist = dict()

def BreakOnAPI():
    global breaklist
    address = pykd.loadDWords(pykd.reg("esp"), 5)
    size = pykd.ptrDWord(address[-1])
    base = (address[2])
    pykd.trace()
    pykd.trace()
    pykd.trace()
    base = pykd.ptrDWord(base)
    if(base != 0):
        bpfs = "ba w4 %x"%base
        breaklist[base] = size
        pykd.dbgCommand(bpfs)
    pykd.go()


def SetZAVMBreak(nt):
    if("ZwAllocateVirtualMemory") in pykd.dbgCommand('bl'):
        return
    b1 = pykd.setBp(nt.ZwAllocateVirtualMemory)
    


def GetMemBase():
    global breaklist
    nt = pykd.loadModule("ntdll")
    SetZAVMBreak(nt)
    pykd.go()
    if(pykd.reg("eip") == nt.ZwAllocateVirtualMemory):
        BreakOnAPI()
        return 0
    for ea,size in breaklist.items():
        cmd = 's -a 0x%X L?0x%X "D030572AD5306501"'%(ea,size)
        if(pykd.dbgCommand(cmd) != ""):
            return ea            
    breaklist = dict()
    pykd.removeBp()
    return 0


def SetMemoryBreakPointer():
    pykd.dbgCommand("sxe ld:flash")
    pykd.go()
    while(1):
        if(GetMemBase()==0):
            continue
        else:
            break

    
SetMemoryBreakPointer()    




